{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use TextGrad with local models and the openai api through LMStudio \n",
    "\n",
    "First, install [LMStudio](https://lmstudio.ai/). Then, use it to download your favorite local model and start a local server. \n",
    "\n",
    "For this demo, we can use one of my favorite models, Maxime Labonne's NeuralBeagle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "from textgrad.engine.local_model_openai_api import ChatExternalClient\n",
    "import textgrad as tg\n",
    "\n",
    "\n",
    "# start a server with lm-studio and point it to the right address; here we use the default address. \n",
    "client = OpenAI(base_url=\"http://localhost:1234/v1\", api_key=\"lm-studio\")\n",
    "\n",
    "engine = ChatExternalClient(client=client, model_string='mlabonne/NeuralBeagle14-7B-GGUF')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it! The rest of the examples will use the local model you downloaded. You just need to pass `engine` by hand. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "tg.set_backward_engine(engine, override=True)\n",
    "\n",
    "initial_solution = \"\"\"To solve the equation 3x^2 - 7x + 2 = 0, we use the quadratic formula:\n",
    "x = (-b ± √(b^2 - 4ac)) / 2a\n",
    "a = 3, b = -7, c = 2\n",
    "x = (7 ± √((-7)^2 + 4(3)(2))) / 6\n",
    "x = (7 ± √73) / 6\n",
    "The solutions are:\n",
    "x1 = (7 + √73)\n",
    "x2 = (7 - √73)\"\"\"\n",
    "\n",
    "solution = tg.Variable(initial_solution,\n",
    "                       requires_grad=True,\n",
    "                       role_description=\"solution to the math question\")\n",
    "\n",
    "loss_system_prompt = tg.Variable(\"\"\"You will evaluate a solution to a math question. \n",
    "Do not attempt to solve it yourself, do not give a solution, only identify errors. Be super concise.\"\"\",\n",
    "                                 requires_grad=False,\n",
    "                                 role_description=\"system prompt\")\n",
    "                              \n",
    "loss_fn = tg.TextLoss(loss_system_prompt)\n",
    "optimizer = tg.TGD([solution])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The given solution is correct and concise. There are no errors identified in the provided evaluation of the quadratic equation. The solutions x1 = (7 + √73) and x2 = (7 - √73) have been calculated correctly using the quadratic formula.\n"
     ]
    }
   ],
   "source": [
    "loss = loss_fn(solution)\n",
    "print(loss.value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "To solve the given quadratic equation 3x^2 - 7x + 2 = 0 using the quadratic formula, we find two solutions:\n",
      "\n",
      "Solution 1: x = (b ± √(b^2 - 4ac)) / 2a = [(7 ± √(7^2 - 4*3*2))] / 6, which simplifies to:\n",
      "   x1 = (7 + √73) / 6\n",
      "\n",
      "Solution 2: Similarly, we calculate the second solution by substituting the negative value of the square root in the equation:\n",
      "   x2 = (7 - √73) / 6\n",
      "\n",
      "It is important to note that these solutions have been calculated using the quadratic formula and are accurate for this specific equation. The quadratic formula, given as (x = (-b ± √(b^2 - 4ac)) / 2a, allows us to find the values of x that satisfy the given equation.\n",
      "<|eot_id|>\n",
      "```\n",
      "\n",
      "```\n",
      "{\n",
      "  \"gradient\": {\n",
      "    \"solutions\": [\n",
      "      {\n",
      "        \"solution_1\": \"x1 = (7 + √73) / 6\"\n",
      "      },\n",
      "      {\n",
      "        \"solution_2\": \"x2 = (7 - √73) / 6\"\n",
      "      },\n",
      "      {\n",
      "        \"explanation\": \"To solve the given quadratic equation 3x^2 - 7x + 2 = 0 using the quadratic formula, we find two solutions: Using the quadratic formula ((x = (-b ± √(b^2 - 4ac)) / 2a) with a = 3, b = -7, and c = 2, we calculate:\n",
      "\n",
      "Solution 1: x = (7 + √73) / 6\n",
      "\n",
      "Solution 2: Similarly, we calculate the second solution by substituting the negative value of the square root in the equation: x = (7 - √73) / 6\n",
      "\n",
      "It is important to note that these solutions have been calculated using the quadratic formula and are accurate for this specific equation. The quadratic formula allows us to find the values of x that satisfy the given equation.\"\n",
      "      }\n",
      "    ]\n",
      "  }\n",
      "}\n",
      "```\n",
      "```\n",
      "\n",
      "<|end_text|>\n",
      "```\n",
      "\n",
      "<|eot_id|>\n",
      "```\n",
      "{\n",
      "  \"gradient\": {\n",
      "    \"solutions\": [\n",
      "      {\n",
      "        \"solution_1\": \"x1 = (7 + √73) / 6\"\n",
      "      },\n",
      "      {\n",
      "        \"solution_2\": \"x2 = (7 - √73) / 6\"\n",
      "      },\n",
      "      {\n",
      "        \"explanation\": \"To solve the given quadratic equation 3x^2 - 7x + 2 = 0 using the quadratic formula, we find two solutions: Using the quadratic formula ((x = (-b ± √(b^2 - 4ac)) / 2a) with a = 3, b = -7, and c = 2, we calculate:\n",
      "\n",
      "Solution 1: x = (7 + √73) / 6\n",
      "\n",
      "Solution 2: Similarly, we calculate the second solution by substituting the negative value of the square root in the equation: x = (7 - √73) / 6\n",
      "\n",
      "It is important to note that these solutions have been calculated using the quadratic formula and are accurate for this specific equation. The quadratic formula allows us to find the values of x that satisfy the given equation.\"\n",
      "      }\n",
      "    ]\n",
      "  }\n",
      "}\n",
      "```\n",
      "```\n",
      "\n",
      "<|end_text|>\n",
      "```\n",
      "\n",
      "The improved variable provides a clearer explanation of how the solutions were calculated using the quadratic formula, along with the specific values for each solution. This enhanced presentation should help readers better understand and follow the process of solving the given equation.\n"
     ]
    }
   ],
   "source": [
    "# this will take a while on CPU. :( \n",
    "loss.backward()\n",
    "optimizer.step()\n",
    "print(solution.value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enjoy passing text gradients through your favorite local models!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
